import numpy as np

class MMSEDetector(object):

    def __init__(self, params):

        self.Nr = params['Nr']
        self.Nt = params['Nt']
        self.batch_size = params['BATCH_SIZE']
        self.C = np.expand_dims(params['CONSTELLATION_COMPLEX'],0) # [1,M]
        self.Rx = np.mean(np.square(np.abs(self.C)))

    def __call__(self,x,H,y,sigma2):
        """
        最小均方误差检测算法
        Input:
        x:发送符号 array(shape=(batch_size, 2*Nt), dtype=float32)
        y:接收信号 array(shape=(batch_size, 2*Nr), dtype=float32)
        H:信道矩阵 array(shape=(batch_size, 2*Nr, 2*Nt), dtype=float32)
        sigma2:噪声方差 array(shape=(batch_size, 1), dtype=float32)
        Output:
        nodes:指标字典
        """
        # Projected channel output
        Hty = self.batch_matvec_mul(np.transpose(H,[0, 2, 1]), y)

        # Gramian of transposed channel matrix
        HtH = np.matmul(np.transpose(H,[0, 2, 1]), H)
        
        # Inverse Gramian
        HtHinv = np.linalg.inv(
            self.Rx * HtH + np.reshape(sigma2/2, [-1,1,1]) * np.eye(2*self.Nt))

        # MMSE Detector
        xhat = self.batch_matvec_mul(HtHinv, self.Rx * Hty)
        
        xhat_complex = self.slicer(xhat)
        x_complex  = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        totalnum = self.batch_size * self.Nt
        rightnum = self.equal(xhat_complex, x_complex)
        errornum = totalnum - rightnum

        nodes={
            'x_complex':x_complex,
            'xhat_complex':xhat_complex,
            'totalnum':totalnum,
            'errornum':errornum,
        }

        return nodes

    def batch_matvec_mul(self, A, b):
        """
        矩阵A与矩阵b相乘，其中A.shape=(batch_size, Nr, Nt)
        b.shape = (batch_size, Nt)
        输出矩阵C，C.shape = (batch_size, Nr)
        """
        C = np.matmul(A, np.expand_dims(b, axis=2))

        return np.squeeze(C, -1)

    def slicer(self, x):
        """
        软符号转换为硬符号
        Input:
        x: 检测器检测后的信号 Tensor(shape=(batchsize, 2*Nt), dtype=float32)
        Output:
        xhat_complex: 符号硬估计，Tensor(shape=(batchsize, Nt)， dtype=tf.int32)
        """
        x_complex = x[:, 0: self.Nt] + 1j*x[:, self.Nt: self.Nt*2]
        x_complex = np.reshape(x_complex,[-1, 1])
        indices = np.argmin(np.abs(x_complex - self.C), axis=1).astype(np.int32)
        xhat_complex = self.C[:,indices].T
        xhat_complex = np.reshape(xhat_complex,[-1,self.Nt])
        
        return xhat_complex

    def equal(self,x,y):
        """
        Computes the numbers of elements for which x and y are equal
        """
        return np.sum(np.equal(x, y).astype(np.float32))